package main

import (
    "encoding/json"
    "flag"
    "fmt"
    "io"
    "io/ioutil"
    "log"
    "net/http"
    "time"

    "golang.org/x/crypto/ssh"
    "gopkg.in/yaml.v3"
)

// Config struct holds settings loaded from a YAML file.
type Config struct {
    JumpHost    string `yaml:"jump_host"`
    JumpUser    string `yaml:"jump_user"`
    JumpPass    string `yaml:"jump_pass"`
    TargetUser  string `yaml:"target_user"`
    TargetPass  string `yaml:"target_pass"`
    InitCommand string `yaml:"init_command"`
    ListenAddr  string `yaml:"listen_addr"`
}

// HookRequest represents the incoming webhook JSON body from wmsshd.
type HookRequest struct {
    Event string `json:"event"`
    Name  string `json:"name"`
    // Add other fields if necessary
}

var (
    config    Config
    debugMode bool
)

func debugLog(format string, args ...interface{}) {
    if debugMode {
        log.Printf("[DEBUG] "+format, args...)
    }
}

func infoLog(format string, args ...interface{}) {
    log.Printf("[INFO] "+format, args...)
}

func errorLog(format string, args ...interface{}) {
    log.Printf("[ERROR] "+format, args...)
}

func printHelp() {
    fmt.Println(`
HookJumpAgent: Automatically initialize new nodes via SSH ProxyJump based on wmsshd webhook event.

Purpose:
    This tool runs an HTTP server to accept post-regist hooks from wmsshd,
    then uses SSH ProxyJump and password authentication (including empty password) 
    to connect to the new node and execute initialization commands.

Command-line arguments:
    -h, --help         Show this help message and exit
    -c, --config       Path to YAML config file (default: config.yaml)
    -d, --debug        Enable debug mode for verbose logging

YAML config example:

    jump_host: wmssh.yuhai.tech:2022
    jump_user: zhangzong
    jump_pass: your_jump_password

    target_user: root
    target_pass: ""             # Empty string means empty password authentication

    init_command: echo "hello world"
    listen_addr: ":8077"

Usage examples:
    ./hookjumpagent -c config.yaml
    ./hookjumpagent -d                      # Enable debug mode
    ./hookjumpagent -c config.yaml -d       # Custom config with debug
    ./hookjumpagent --help

Doc & feedback:
    https://your-company-docs-url/example

`)
}

func main() {
    showHelp := flag.Bool("help", false, "Show help")
    showHelp2 := flag.Bool("h", false, "Show help")
    configFile := flag.String("config", "config.yaml", "YAML config file")
    configFile2 := flag.String("c", "config.yaml", "YAML config file")
    debug := flag.Bool("debug", false, "Enable debug mode")
    debug2 := flag.Bool("d", false, "Enable debug mode")
    flag.Parse()

    if *showHelp || *showHelp2 {
        printHelp()
        return
    }

    debugMode = *debug || *debug2

    cfgFile := *configFile
    if flag.Lookup("c").Value.String() != "config.yaml" {
        cfgFile = *configFile2
    }

    infoLog("Starting HookJumpAgent...")
    debugLog("Loading config file: %s", cfgFile)
    debugLog("Debug mode enabled")

    data, err := ioutil.ReadFile(cfgFile)
    if err != nil {
        errorLog("Failed to read config file '%s': %v", cfgFile, err)
        log.Fatalf("Failed to read config: %v", err)
    }

    if err := yaml.Unmarshal(data, &config); err != nil {
        errorLog("Failed to parse YAML config: %v", err)
        log.Fatalf("Failed to parse yaml: %v", err)
    }

    if config.ListenAddr == "" {
        config.ListenAddr = ":8077"
        debugLog("Using default listen address: %s", config.ListenAddr)
    }

    debugLog("Config loaded successfully:")
    debugLog("  Jump Host: %s", config.JumpHost)
    debugLog("  Jump User: %s", config.JumpUser)
    debugLog("  Jump Pass: %s", maskPassword(config.JumpPass))
    debugLog("  Target User: %s", config.TargetUser)
    debugLog("  Target Pass: %s", maskPassword(config.TargetPass))
    debugLog("  Init Command: %s", config.InitCommand)
    debugLog("  Listen Address: %s", config.ListenAddr)

    http.HandleFunc("/report", HookHandler)
    infoLog("HookJumpAgent started successfully, listening on %s", config.ListenAddr)
    
    if debugMode {
        infoLog("Debug mode is enabled - verbose logging is active")
    }
    
    errorLog("HTTP server stopped: %v", http.ListenAndServe(config.ListenAddr, nil))
}

func maskPassword(password string) string {
    if password == "" {
        return "<empty>"
    }
    if len(password) <= 2 {
        return "***"
    }
    return password[:1] + "***" + password[len(password)-1:]
}

func HookHandler(w http.ResponseWriter, r *http.Request) {
    debugLog("Received webhook request from %s", r.RemoteAddr)
    debugLog("Request method: %s, URL: %s", r.Method, r.URL.Path)
    
    defer r.Body.Close()
    body, err := io.ReadAll(r.Body)
    if err != nil {
        errorLog("Failed to read request body: %v", err)
        w.WriteHeader(http.StatusBadRequest)
        return
    }

    debugLog("Request body: %s", string(body))

    var req HookRequest
    if err := json.Unmarshal(body, &req); err != nil {
        errorLog("Invalid JSON in webhook request: %v", err)
        debugLog("Raw body that failed to parse: %s", string(body))
        w.WriteHeader(http.StatusOK)
        return
    }

    debugLog("Parsed webhook request - Event: %s, Name: %s", req.Event, req.Name)

    if req.Event == "post-regist" {
        node := req.Name
        infoLog("New node registration detected: %s", node)
        go func(node string) {
            debugLog("Starting initialization goroutine for node: %s", node)
            err := sshJumpAndInit(node)
            if err != nil {
                errorLog("Failed to initialize node %s: %v", node, err)
            } else {
                infoLog("Successfully initialized node: %s", node)
            }
        }(node)
    } else {
        debugLog("Ignoring webhook event: %s (not post-regist)", req.Event)
    }

    w.WriteHeader(http.StatusOK)
    w.Write([]byte("ok"))
    debugLog("Webhook request processed successfully")
}

func sshJumpAndInit(targetHost string) error {
    infoLog("Initializing node %s via SSH jump", targetHost)
    
    // 1. Connect to jump host using password
    debugLog("Connecting to jump host: %s", config.JumpHost)
    jumpConfig := &ssh.ClientConfig{
        User:            config.JumpUser,
        Auth:            []ssh.AuthMethod{ssh.Password(config.JumpPass)},
        HostKeyCallback: ssh.InsecureIgnoreHostKey(),
        Timeout:         10 * time.Second,
    }
    
    jumpClient, err := ssh.Dial("tcp", config.JumpHost, jumpConfig)
    if err != nil {
        errorLog("Failed to connect to jump host %s: %v", config.JumpHost, err)
        return fmt.Errorf("Failed to dial jump host: %v", err)
    }
    defer jumpClient.Close()
    debugLog("Successfully connected to jump host")

    // 2. Connect from jump host to target node
    taddr := targetHost
    if len(taddr) > 0 && taddr[len(taddr)-3:] != ":22" {
        taddr += ":22"
    }
    
    debugLog("Connecting to target host %s via jump host", taddr)
    conn, err := jumpClient.Dial("tcp", taddr)
    if err != nil {
        errorLog("Failed to connect to target %s via jump host: %v", taddr, err)
        return fmt.Errorf("Failed to dial target via jump: %v", err)
    }
    debugLog("Successfully established connection to target host via jump")

    // 3. Prepare target node ssh client config (always using password, even if empty)
    debugLog("Authenticating to target host as user: %s", config.TargetUser)
    targetConfig := &ssh.ClientConfig{
        User:            config.TargetUser,
        Auth:            []ssh.AuthMethod{ssh.Password(config.TargetPass)}, // will use empty password if config.TargetPass == ""
        HostKeyCallback: ssh.InsecureIgnoreHostKey(),
        Timeout:         10 * time.Second,
    }

    clientConn, chans, reqs, err := ssh.NewClientConn(conn, taddr, targetConfig)
    if err != nil {
        errorLog("Failed to authenticate to target host %s: %v", taddr, err)
        return fmt.Errorf("Failed to build new client conn: %v", err)
    }
    targetClient := ssh.NewClient(clientConn, chans, reqs)
    defer targetClient.Close()
    debugLog("Successfully authenticated to target host")

    // 4. Execute initialization command on target node
    debugLog("Executing initialization command: %s", config.InitCommand)
    session, err := targetClient.NewSession()
    if err != nil {
        errorLog("Failed to create SSH session: %v", err)
        return fmt.Errorf("Failed to create session: %v", err)
    }
    defer session.Close()
    
    startTime := time.Now()
    output, err := session.CombinedOutput(config.InitCommand)
    duration := time.Since(startTime)
    
    if err != nil {
        errorLog("Command execution failed on %s (took %v): %v", targetHost, duration, err)
        debugLog("Command output: %s", string(output))
        return fmt.Errorf("Failed to run cmd: %v (output=%s)", err, string(output))
    }
    
    infoLog("Command executed successfully on %s (took %v)", targetHost, duration)
    if debugMode {
        debugLog("Command output from %s:\n%s", targetHost, string(output))
    } else {
        infoLog("Command output from %s: %d bytes", targetHost, len(output))
    }
    
    return nil
}


